import torch
from torch.utils.data import DataLoader
from torchvision import datasets,transforms
from  torch import nn,optim

import ssl
ssl._create_default_https_context = ssl._create_unverified_context

batch_size=1024
mnist_train=datasets.MNIST("mnist",True,transform=transforms.Compose([
    transforms.ToTensor()    ]),download=True)
mnist_train=DataLoader(mnist_train,batch_size=batch_size,shuffle=True)
minst_test=datasets.MNIST("mnist",False,transform=transforms.Compose([
    transforms.ToTensor()  ]),download=True)
minst_test=DataLoader(minst_test,batch_size=batch_size,shuffle=True)
x,lable=next(iter(mnist_train))
print(lable)
x.shape

device=torch.device("mps")
autoencoder=AE().to(device)
critenon=nn.MSELoss()
optimizer=optim.Adam(autoencoder.parameters(),lr=1e-4)

autoencoder2=AE()
critenon2=nn.MSELoss()
optimizer2=optim.Adam(autoencoder2.parameters(),lr=1e-4)

# GPU 训练
#%%time
for epoch in range(5):
    for index,(x,_) in enumerate(mnist_train):
        x=x.to(device)
        x_hat=autoencoder(x)
        loss=critenon(x_hat,x)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
    print(epoch,"loss: ",loss.item())
    
# CPU训练
# %%time
for epoch in range(5):
    for index,(x,_) in enumerate(mnist_train):
        x=x
        x_hat=autoencoder2(x)
        loss=critenon2(x_hat,x)
        optimizer2.zero_grad()
        loss.backward()
        optimizer2.step()
    print(epoch,"loss: ",loss.item())

total_params = sum(p.numel() for p in autoencoder2.parameters())
print("Total Parameters: {:,}".format(total_params))
